{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Long Short-Term Memory](https://dl.acm.org/doi/10.1162/neco.1997.9.8.1735)\n",
    "\n",
    "\n",
    "[1997-LSTM.pdf](../papers/1997-LSTM.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![LSTM示意图](http://assets.hypervoid.top/img/2025/06/30/image-20250630170007753-a624.png)\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "i_t &= \\sigma(W_{ii}x_t + b_{ii} + W_{hi}h_{t-1} + b_{hi}) \\\\\n",
    "f_t &= \\sigma(W_{if}x_t + b_{if} + W_{hf}h_{t-1} + b_{hf}) \\\\\n",
    "g_t &= \\tanh(W_{ig}x_t + b_{ig} + W_{hg}h_{t-1} + b_{hg}) \\\\\n",
    "o_t &= \\sigma(W_{io}x_t + b_{io} + W_{ho}h_{t-1} + b_{ho}) \\\\\n",
    "c_t &= f_t \\odot c_{t-1} + i_t \\odot g_t \\\\\n",
    "h_t &= o_t \\odot \\tanh(c_t)\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "torch.manual_seed(20250630)\n",
    "torch.cuda.manual_seed(20250630)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size, seq_len = 2, 3\n",
    "# 输入大小，隐含层大小\n",
    "input_size, hidden_size = 4, 5\n",
    "# 输入\n",
    "input = torch.randn(batch_size, seq_len, input_size)\n",
    "# 初始值\n",
    "c0, h0 = torch.randn(batch_size, hidden_size), torch.randn(batch_size, hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.4276,  0.2803,  0.0205, -0.0904, -0.0928],\n",
      "         [ 0.3246,  0.0375,  0.1131, -0.0302, -0.4382],\n",
      "         [ 0.2796, -0.2374,  0.1253, -0.0093, -0.2690]],\n",
      "\n",
      "        [[ 0.3898, -0.1509,  0.0402,  0.0404,  0.3354],\n",
      "         [ 0.2717, -0.2599,  0.1875, -0.0164,  0.3097],\n",
      "         [ 0.2790, -0.4573,  0.1867,  0.0285,  0.3013]]],\n",
      "       grad_fn=<TransposeBackward0>)\n",
      "tensor([[[ 0.2796, -0.2374,  0.1253, -0.0093, -0.2690],\n",
      "         [ 0.2790, -0.4573,  0.1867,  0.0285,  0.3013]]],\n",
      "       grad_fn=<StackBackward0>)\n",
      "tensor([[[ 0.6250, -0.4408,  0.2818, -0.0264, -0.4600],\n",
      "         [ 0.7244, -0.9398,  0.4891,  0.1022,  0.4883]]],\n",
      "       grad_fn=<StackBackward0>)\n",
      "---Parameters---\n",
      "weight_ih_l0:\ttorch.Size([20, 4])\n",
      "weight_hh_l0:\ttorch.Size([20, 5])\n",
      "bias_ih_l0:\ttorch.Size([20])\n",
      "bias_hh_l0:\ttorch.Size([20])\n"
     ]
    }
   ],
   "source": [
    "# Pytorch LSTM\n",
    "torch_lstm = nn.LSTM(input_size, hidden_size, batch_first=True)\n",
    "out_torch, (h_torch, c_torch) = torch_lstm.forward(\n",
    "    input,\n",
    "    (h0.unsqueeze(0), c0.unsqueeze(0)),\n",
    ")\n",
    "print(out_torch)\n",
    "print(h_torch)\n",
    "print(c_torch)\n",
    "print(\"---Parameters---\")\n",
    "for k, v in torch_lstm.named_parameters():\n",
    "    print(f\"{k}:\\t{v.shape}\")\n",
    "# hidden_size=5, ifgo 一共四个w拼起来 就是 20\n",
    "# weight_ih_l0:\ttorch.Size([20, 4]) # \n",
    "# weight_hh_l0:\ttorch.Size([20, 5]) # \n",
    "# bias_ih_l0:\ttorch.Size([20])    # \n",
    "# bias_hh_l0:\ttorch.Size([20])    # "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "from torch import Tensor\n",
    "\n",
    "\n",
    "class LstmCell(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super(LstmCell, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.w_i = nn.Parameter(torch.randn(4 * hidden_size, input_size))\n",
    "        self.w_h = nn.Parameter(torch.randn(4 * hidden_size, hidden_size))\n",
    "        self.b_i = nn.Parameter(torch.randn(4 * hidden_size))\n",
    "        self.b_h = nn.Parameter(torch.randn(4 * hidden_size))\n",
    "        self.init_weights()\n",
    "\n",
    "    def init_weights(self):\n",
    "        nn.init.xavier_uniform_(self.w_i)\n",
    "        nn.init.xavier_uniform_(self.w_h)\n",
    "        nn.init.zeros_(self.b_i)\n",
    "        nn.init.zeros_(self.b_h)\n",
    "\n",
    "    def forward(\n",
    "        self, x_t: Tensor, state: tuple[Tensor, Tensor]\n",
    "    ) -> tuple[Tensor, Tensor]:\n",
    "        h_prev, c_prev = state\n",
    "        gates = ((x_t @ self.w_i.T) + (h_prev @ self.w_h.T) + self.b_i + self.b_h)  # [1, 2, 20]\n",
    "        gates.squeeze_(0)\n",
    "        # print(gates.shape)\n",
    "        input_gate, forget_gate, cell_gate, output_gate = gates.chunk(4, 1)\n",
    "\n",
    "        # 分别计算 输入门(i)、遗忘门(f)、cell门(f)、输出门(o)\n",
    "        i_t = torch.sigmoid(input_gate)\n",
    "        f_t = torch.sigmoid(forget_gate)\n",
    "        g_t = torch.tanh(cell_gate)\n",
    "        o_t = torch.sigmoid(output_gate)\n",
    "        c_next = f_t * c_prev + i_t * g_t\n",
    "        h_next = o_t * torch.tanh(c_next)\n",
    "        return h_next, c_next\n",
    "\n",
    "\n",
    "class LSTM(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super(LSTM, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.lstm_cell = LstmCell(input_size, hidden_size)\n",
    "        self.leaky_relu = nn.LeakyReLU()\n",
    "\n",
    "    def forward(self, x: Tensor, init: Tuple[Tensor, Tensor] | None = None):\n",
    "        batch_size_, seq_len_, _ = x.shape\n",
    "        h_out = torch.zeros((batch_size_, seq_len_, self.hidden_size))\n",
    "        if init is None:\n",
    "            h_prev = torch.zeros(batch_size, self.hidden_size, device=x.device)\n",
    "            c_prev = torch.zeros(batch_size, self.hidden_size, device=x.device)\n",
    "        else:\n",
    "            h_prev, c_prev = init\n",
    "\n",
    "        for i in range(seq_len):\n",
    "            x_t = x[:, i, :]\n",
    "            h_prev, c_prev = self.lstm_cell(x_t, (h_prev, c_prev))\n",
    "            h_out[:, i, :] = h_prev\n",
    "        return h_out, (h_prev, c_prev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align*}\n",
    "i_t &= \\sigma(W_{ii}x_t + b_{ii} + W_{hi}h_{t-1} + b_{hi}) \\\\\n",
    "f_t &= \\sigma(W_{if}x_t + b_{if} + W_{hf}h_{t-1} + b_{hf}) \\\\\n",
    "g_t &= \\tanh(W_{ig}x_t + b_{ig} + W_{hg}h_{t-1} + b_{hg}) \\\\\n",
    "o_t &= \\sigma(W_{io}x_t + b_{io} + W_{ho}h_{t-1} + b_{ho}) \\\\\n",
    "c_t &= f_t \\odot c_{t-1} + i_t \\odot g_t \\\\\n",
    "h_t &= o_t \\odot \\tanh(c_t)\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.4276,  0.2803,  0.0205, -0.0904, -0.0928],\n",
      "         [ 0.3246,  0.0375,  0.1131, -0.0302, -0.4382],\n",
      "         [ 0.2796, -0.2374,  0.1253, -0.0093, -0.2690]],\n",
      "\n",
      "        [[ 0.3898, -0.1509,  0.0402,  0.0404,  0.3354],\n",
      "         [ 0.2717, -0.2599,  0.1875, -0.0164,  0.3097],\n",
      "         [ 0.2790, -0.4573,  0.1867,  0.0285,  0.3013]]], grad_fn=<CopySlices>)\n",
      "tensor([[[ 0.2796, -0.2374,  0.1253, -0.0093, -0.2690],\n",
      "         [ 0.2790, -0.4573,  0.1867,  0.0285,  0.3013]]],\n",
      "       grad_fn=<MulBackward0>)\n",
      "tensor([[[ 0.6250, -0.4408,  0.2818, -0.0264, -0.4600],\n",
      "         [ 0.7244, -0.9398,  0.4891,  0.1022,  0.4883]]],\n",
      "       grad_fn=<AddBackward0>)\n",
      "---Parameters---\n",
      "lstm_cell.w_i:\ttorch.Size([20, 4])\n",
      "lstm_cell.w_h:\ttorch.Size([20, 5])\n",
      "lstm_cell.b_i:\ttorch.Size([20])\n",
      "lstm_cell.b_h:\ttorch.Size([20])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# Pytorch LSTM\n",
    "my_lstm = LSTM(input_size, hidden_size)\n",
    "my_lstm.lstm_cell.w_i = torch_lstm.weight_ih_l0\n",
    "my_lstm.lstm_cell.w_h = torch_lstm.weight_hh_l0\n",
    "my_lstm.lstm_cell.b_i = torch_lstm.bias_ih_l0\n",
    "my_lstm.lstm_cell.b_h = torch_lstm.bias_hh_l0\n",
    "out_, (h_, c_) = my_lstm.forward(\n",
    "    input,\n",
    "    (h0.unsqueeze(0), c0.unsqueeze(0)),\n",
    ")\n",
    "\n",
    "\n",
    "\n",
    "print(out_)\n",
    "print(h_)\n",
    "print(c_)\n",
    "print(\"---Parameters---\")\n",
    "for k, v in my_lstm.named_parameters():\n",
    "    print(f\"{k}:\\t{v.shape}\")\n",
    "# tensor([[[-0.0310, -0.0355,  0.2224,  0.2715, -0.2373],\n",
    "#          [-0.0458, -0.0054,  0.2440,  0.0589, -0.5580],\n",
    "#          [ 0.0597, -0.2669,  0.1833,  0.0118, -0.3843]],\n",
    "\n",
    "#         [[ 0.1301, -0.0079,  0.1585, -0.2035, -0.1400],\n",
    "#          [ 0.1941, -0.2171,  0.1993, -0.0618, -0.0919],\n",
    "#          [ 0.2392, -0.4258,  0.1861,  0.0149,  0.1375]]],\n",
    "#        grad_fn=<TransposeBackward0>)\n",
    "# tensor([[[ 0.0597, -0.2669,  0.1833,  0.0118, -0.3843],\n",
    "#          [ 0.2392, -0.4258,  0.1861,  0.0149,  0.1375]]],\n",
    "#        grad_fn=<StackBackward0>)\n",
    "# tensor([[[ 0.1169, -0.5324,  0.4434,  0.0319, -0.6892],\n",
    "#          [ 0.5470, -0.8999,  0.4868,  0.0556,  0.2197]]],"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert c_.allclose(c_torch)\n",
    "assert h_.allclose(h_torch)\n",
    "assert out_.allclose(out_torch)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
